"""
@Description :   测试器基类
@Author      :   tqychy 
@Time        :   2025/08/28 08:57:42
"""

import os

import torchvision.transforms as T
from PIL import Image
from torch.utils.data import DataLoader

from dataset import datasets
from metrics import metrics_handlers
from nets import models


class BaseTester:
    def __init__(self, *args):
        self.cfg, self.logger = args
        self.result_path = os.path.join("results", self.cfg.GLOBALS.EXPR_NAME)
        self.logger.info(self.cfg)
        # 加载测试数据集
        dataset_name = self.cfg.EVAL.DATASET
        max_len = self.cfg.EVAL.DATASET_MAX_LEN
        self.dataset = datasets[dataset_name](
            self.cfg, self.logger, dataset=dataset_name, split="val", max_len=max_len)
        self.dataloader = DataLoader(self.dataset, batch_size=1, shuffle=False)
        # 加载模型
        model_path = self.cfg.EVAL.MODEL_PATH
        self.model = models[self.cfg.EVAL.MODEL_NAME](
            self.cfg, self.logger, model_path=model_path)
        # 构建结果保存目录
        self.pred_save_path = os.path.join(self.result_path, "pred_bbox")
        self.gt_save_path = os.path.join(self.result_path, "gt_bbox")
        os.makedirs(self.pred_save_path, exist_ok=True)
        os.makedirs(self.gt_save_path, exist_ok=True)
        # 构建指标处理 Handler
        self.metrics_handler = metrics_handlers[self.cfg.EVAL.TYPE](self.cfg, self.logger)

    @staticmethod
    def load_image(image_path):
        transforms = T.ToTensor()
        image = Image.open(image_path).convert("RGB")
        image = transforms(image)
        image = image.squeeze().permute(1, 2, 0)
        return image

    def covert_formatted_bbox(self, bbox, image_shape):
        """
        根据模型输出将不同格式的 bbox 转换为 [x_min, y_min, width, height] 格式

        """
        return self.model.covert_formatted_bbox(bbox, image_shape)

    def make_prompt(self, image_path, sentence, category):
        return self.model.make_prompt(image_path, sentence, category, self.cfg.EVAL.TYPE)

    @staticmethod
    def parse_output(output: str):
        raise NotImplementedError("parse_output 方法未实现")

    def test(self):
        raise NotImplementedError("test 方法未实现")
